#pragma once

#include "updater_common.h"

namespace mdtx::online::updater
{

    // TODO: enable Kahan compensation for arma vectors

    // Covariance matrix, naiive impl
    struct cov_updater
    {
        size_t n;
        arma::rowvec m;
        arma::mat s2;

        explicit cov_updater(int nvar)
            : n(0),
              m(nvar, arma::fill::zeros),
              s2(nvar, nvar, arma::fill::zeros)
        {
        }

        void insert(const arma::rowvec &x)
        {
            arma::rowvec d{x - m};
            n += 1;
            m += d / n;
            s2 += (x - m).t() * d;
        }

        void remove(const arma::rowvec &x)
        {
            arma::rowvec d{x - m};
            n -= 1;
            m -= d / n;
            s2 -= (x - m).t() * d;
        }
    };

    // Exponentailly weighted covariance matrix, naiive impl
    struct ewcov_updater
    {
        bool init;
        double a, b;
        arma::rowvec m;
        arma::mat s2;

        ewcov_updater(int nvar, double alpha)
            : init(false),
              a(alpha),
              b(1 - a),
              m(nvar, arma::fill::zeros),
              s2(nvar, nvar, arma::fill::zeros)
        {
        }

        ewcov_updater(int nvar, int period)
            : ewcov_updater(nvar, 2.0 / (period + 1))
        {
        }

        void insert(const arma::rowvec &x)
        {
            if (init)
            {
                arma::rowvec d{x - m};
                arma::colvec i{a * d};
                m += i;
                s2 += i * d;
                s2 *= b;
            }
            else
            {
                m = x;
                init = true;
            }
        }
    };

    // arithmetic
    struct s2v_updater
    {
        size_t n;
        arma::rowvec m, s2;

        explicit s2v_updater(int nvar)
            : n(0),
              m(nvar, arma::fill::zeros),
              s2(nvar, arma::fill::zeros)
        {
        }

        void insert(const arma::rowvec &x)
        {
            arma::rowvec d{x - m};
            n += 1;
            m += d / n;
            s2 += (x - m) % d;
        }

        void remove(const arma::rowvec &x)
        {
            arma::rowvec d{x - m};
            n -= 1;
            m -= d / n;
            s2 -= (x - m) % d;
        }
    };

    // weighted
    struct s2wv_updater
    {
        arma::rowvec wsum, m, s2;

        explicit s2wv_updater(int nvar)
            : wsum(nvar, arma::fill::zeros),
              m(nvar, arma::fill::zeros),
              s2(nvar, arma::fill::zeros)
        {
        }

        void insert(const arma::rowvec &x, const arma::rowvec &w)
        {
            arma::rowvec d = x - m;
            wsum += w;
            m += (w / wsum) % d;
            s2 += w % (x - m) % d;
        }

        void remove(const arma::rowvec &x, const arma::rowvec &w)
        {
            arma::rowvec d = x - m;
            wsum -= w;
            m -= (w / wsum) % d;
            s2 -= w % (x - m) % d;
        }
    };

    // weighted, w^2 is also calculated
    struct s2wrv_updater
    {
        arma::rowvec wsum, w2sum, m, s2;

        explicit s2wrv_updater(int nvar)
            : wsum(nvar, arma::fill::zeros),
              w2sum(nvar, arma::fill::zeros),
              m(nvar, arma::fill::zeros),
              s2(nvar, arma::fill::zeros)
        {
        }

        void insert(const arma::rowvec &x, const arma::rowvec &w)
        {
            arma::rowvec d = x - m;
            wsum += w;
            w2sum += w % w;
            m += (w / wsum) % d;
            s2 += w % (x - m) % d;
        }

        void remove(const arma::rowvec &x, const arma::rowvec &w)
        {
            arma::rowvec d = x - m;
            wsum -= w;
            w2sum -= w % w;
            m -= (w / wsum) % d;
            s2 -= w % (x - m) % d;
        }
    };

}
